#
from typing import Dict
import matplotlib.pyplot as plt
from mmlm.ce.expand_context import ExpandContext

class ExpandContextApp(object):
    def __init__(self):
        self.name = 'apps.ces.expand_context_app.ExpandContextApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'扩展上下文例程......v0.0.2')
        if params['task'] == 1:
            ExpandContextApp.context_expansion_exp()
        elif params['task'] == 2:
            ExpandContextApp.test_context_expand_template(params=params)
        elif params['task'] == 3:
            ExpandContextApp.test_layered_contexts(params=params)
        elif params['task'] == 4:
            ExpandContextApp.test_compress_context(params=params)
        elif params['task'] == 5:
            ExpandContextApp.test_prune_context_layers(params=params)
        elif params['task'] == 6:
            ExpandContextApp.test_context_expand_retrieval(params=params)

    def context_expansion_exp() -> None:
        # Base prompt (atom)
        base_prompt = "Write a paragraph about climate change."
        # Expanded prompt variations (molecules)
        expanded_prompts = {
            "base": base_prompt,
            
            "with_role": """You are an environmental scientist with expertise in climate systems. 
        Write a paragraph about climate change.""",
            
            "with_examples": """Write a paragraph about climate change.

        Example 1:
        Climate change refers to long-term shifts in temperatures and weather patterns. Human activities have been the main driver of climate change since the 1800s, primarily due to the burning of fossil fuels like coal, oil, and gas, which produces heat-trapping gases.

        Example 2:
        Global climate change is evident in the increasing frequency of extreme weather events, rising sea levels, and shifting wildlife populations. Scientific consensus points to human activity as the primary cause.""",
            
            "with_constraints": """Write a paragraph about climate change.
        - Include at least one scientific fact with numbers
        - Mention both causes and effects
        - End with a call to action
        - Keep the tone informative but accessible""",
            
            "with_audience": """Write a paragraph about climate change for high school students who are
        just beginning to learn about environmental science. Use clear explanations 
        and relatable examples.""",
            
            "comprehensive": """You are an environmental scientist with expertise in climate systems.

        Write a paragraph about climate change for high school students who are
        just beginning to learn about environmental science. Use clear explanations 
        and relatable examples.

        Guidelines:
        - Include at least one scientific fact with numbers
        - Mention both causes and effects
        - End with a call to action
        - Keep the tone informative but accessible

        Example of tone and structure:
        "Ocean acidification occurs when seawater absorbs CO2 from the atmosphere, causing pH levels to drop. Since the Industrial Revolution, ocean pH has decreased by 0.1 units, representing a 30% increase in acidity. This affects marine life, particularly shellfish and coral reefs, as it impairs their ability to form shells and skeletons. Scientists predict that if emissions continue at current rates, ocean acidity could increase by 150% by 2100, devastating marine ecosystems. By reducing our carbon footprint through simple actions like using public transportation, we can help protect these vital ocean habitats."
        """
        }
        results = {}
        responses = {}

        for name, prompt in expanded_prompts.items():
            print(f"Testing prompt: {name}")
            response, latency = ExpandContext.generate_response(prompt)
            responses[name] = response
            metrics = ExpandContext.calculate_metrics(prompt, response, latency)
            results[name] = metrics
            print(f"  Prompt tokens: {metrics['prompt_tokens']}")
            print(f"  Response tokens: {metrics['response_tokens']}")
            print(f"  Latency: {metrics['latency']:.2f}s")
            print("-" * 40)

        # Prepare data for visualization
        prompt_types = list(results.keys())
        prompt_tokens = [results[k]['prompt_tokens'] for k in prompt_types]
        response_tokens = [results[k]['response_tokens'] for k in prompt_types]
        latencies = [results[k]['latency'] for k in prompt_types]

        # Create figure with multiple subplots
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))

        # Plot 1: Token Usage
        axes[0, 0].bar(prompt_types, prompt_tokens, label='Prompt Tokens', alpha=0.7, color='blue')
        axes[0, 0].bar(prompt_types, response_tokens, bottom=prompt_tokens, label='Response Tokens', alpha=0.7, color='green')
        axes[0, 0].set_title('Token Usage by Prompt Type')
        axes[0, 0].set_ylabel('Number of Tokens')
        axes[0, 0].legend()
        plt.setp(axes[0, 0].get_xticklabels(), rotation=45, ha='right')

        # Plot 2: Token Efficiency (Response Tokens / Prompt Tokens)
        token_efficiency = [results[k]['token_efficiency'] for k in prompt_types]
        axes[0, 1].bar(prompt_types, token_efficiency, color='purple', alpha=0.7)
        axes[0, 1].set_title('Token Efficiency (Response/Prompt)')
        axes[0, 1].set_ylabel('Efficiency Ratio')
        plt.setp(axes[0, 1].get_xticklabels(), rotation=45, ha='right')

        # Plot 3: Latency
        axes[1, 0].bar(prompt_types, latencies, color='red', alpha=0.7)
        axes[1, 0].set_title('Response Latency')
        axes[1, 0].set_ylabel('Seconds')
        plt.setp(axes[1, 0].get_xticklabels(), rotation=45, ha='right')

        # Plot 4: Latency per 1k tokens
        latency_per_1k = [results[k]['latency_per_1k'] for k in prompt_types]
        axes[1, 1].bar(prompt_types, latency_per_1k, color='orange', alpha=0.7)
        axes[1, 1].set_title('Latency per 1k Tokens')
        axes[1, 1].set_ylabel('Seconds per 1k Tokens')
        plt.setp(axes[1, 1].get_xticklabels(), rotation=45, ha='right')

        plt.tight_layout()
        plt.show()

        # 查看响应内容，用于评估响应的质量
        for name, response in responses.items():
            print(f"=== Response for {name} prompt ===")
            print(response)
            print("\n" + "=" * 80 + "\n")

    @staticmethod
    def test_context_expand_template(params:Dict = {}) -> None:
        # Test our template
        new_base_prompt = "Explain how photosynthesis works."

        new_expanded_context = ExpandContext.create_expanded_context(
            base_prompt=new_base_prompt,
            role="a biology teacher with 15 years of experience",
            audience="middle school students",
            tone="enthusiastic and educational",
            constraints=[
                "Use a plant-to-factory analogy",
                "Mention the role of chlorophyll",
                "Explain the importance for Earth's ecosystem",
                "Keep it under 200 words"
            ],
            examples=[
                "Photosynthesis is like a tiny factory inside plants. Just as a factory needs raw materials, energy, and workers to make products, plants need carbon dioxide, water, sunlight, and chlorophyll to make glucose (sugar) and oxygen. The sunlight is the energy source, chlorophyll molecules are the workers that capture this energy, while carbon dioxide and water are the raw materials. The factory's products are glucose, which the plant uses for growth and energy storage, and oxygen, which is released into the air for animals like us to breathe. This process is essential for life on Earth because it provides the oxygen we need and removes carbon dioxide from the atmosphere."
            ]
        )

        print("Template-generated expanded context:")
        print("-" * 80)
        print(new_expanded_context)
        print("-" * 80)
        print(f"Token count: {ExpandContext.count_tokens(new_expanded_context)}")

        # Generate a response using our expanded context
        response, latency = ExpandContext.generate_response(new_expanded_context)
        metrics = ExpandContext.calculate_metrics(new_expanded_context, response, latency)

        print("\nResponse:")
        print("-" * 80)
        print(response)
        print("-" * 80)
        print(f"Response tokens: {metrics['response_tokens']}")
        print(f"Latency: {metrics['latency']:.2f}s")

    @staticmethod
    def test_layered_contexts(params:Dict = {}) -> None:
        layer_test_prompt = "Write code to implement a simple weather app."

        context_layers = {
            "role": "You are a senior software engineer with expertise in full-stack development and UI/UX design.",
            
            "requirements": """Requirements:
        - The app should show current temperature, conditions, and forecast for the next 3 days
        - It should allow users to search for weather by city name
        - It should have a clean, responsive interface
        - The app should handle error states gracefully""",
            
            "tech_stack": """Technical specifications:
        - Use HTML, CSS, and vanilla JavaScript (no frameworks)
        - Use the OpenWeatherMap API for weather data
        - All code should be well-commented and follow best practices
        - Include both the HTML structure and JavaScript functionality""",
            
            "example": """Example structure (but improve upon this):
        ```html
        <!DOCTYPE html>
        <html>
        <head>
            <title>Weather App</title>
            <link rel="stylesheet" href="styles.css">
        </head>
        <body>
            <div class="container">
                <h1>Weather App</h1>
                <div class="search">
                    <input type="text" placeholder="Enter city name">
                    <button>Search</button>
                </div>
                <div class="weather-display">
                    <!-- Weather data will be displayed here -->
                </div>
            </div>
            <script src="app.js"></script>
        </body>
        </html>
        ```"""
        }

        # Run the layer optimization test
        layer_test_results = ExpandContext.layered_contexts(layer_test_prompt, context_layers)
        # Extract data for visualization
        config_names = list(layer_test_results.keys())
        prompt_sizes = [layer_test_results[k]['prompt_tokens'] for k in config_names]
        response_sizes = [layer_test_results[k]['response_tokens'] for k in config_names]
        efficiencies = [layer_test_results[k]['token_efficiency'] for k in config_names]

        # Create visualization
        fig, axes = plt.subplots(2, 1, figsize=(12, 10))

        # Plot 1: Token usage by configuration
        axes[0].bar(config_names, prompt_sizes, label='Prompt Tokens', alpha=0.7, color='blue')
        axes[0].bar(config_names, response_sizes, label='Response Tokens', alpha=0.7, color='green')
        axes[0].set_title('Token Usage by Context Configuration')
        axes[0].set_ylabel('Number of Tokens')
        axes[0].legend()
        plt.setp(axes[0].get_xticklabels(), rotation=45, ha='right')

        # Plot 2: Token efficiency by configuration
        axes[1].bar(config_names, efficiencies, color='purple', alpha=0.7)
        axes[1].set_title('Token Efficiency by Context Configuration')
        axes[1].set_ylabel('Efficiency Ratio (Response/Prompt)')
        plt.setp(axes[1].get_xticklabels(), rotation=45, ha='right')

        plt.tight_layout()
        plt.show()

        # Identify the most efficient configuration
        most_efficient = max(config_names, key=lambda x: layer_test_results[x]['token_efficiency'])
        print(f"Most token-efficient configuration: {most_efficient}")
        print(f"Efficiency ratio: {layer_test_results[most_efficient]['token_efficiency']:.2f}")

    @staticmethod
    def test_compress_context(params:Dict = {}) -> None:
        new_base_prompt = "Explain how photosynthesis works."

        original_context = ExpandContext.create_expanded_context(
            base_prompt=new_base_prompt,
            role="a biology teacher with 15 years of experience",
            audience="middle school students",
            tone="enthusiastic and educational",
            constraints=[
                "Use a plant-to-factory analogy",
                "Mention the role of chlorophyll",
                "Explain the importance for Earth's ecosystem",
                "Keep it under 200 words"
            ],
            examples=[
                "Photosynthesis is like a tiny factory inside plants. Just as a factory needs raw materials, energy, and workers to make products, plants need carbon dioxide, water, sunlight, and chlorophyll to make glucose (sugar) and oxygen. The sunlight is the energy source, chlorophyll molecules are the workers that capture this energy, while carbon dioxide and water are the raw materials. The factory's products are glucose, which the plant uses for growth and energy storage, and oxygen, which is released into the air for animals like us to breathe. This process is essential for life on Earth because it provides the oxygen we need and removes carbon dioxide from the atmosphere."
            ]
        )
        print(f"Original context token count: {ExpandContext.count_tokens(original_context)}")

        for technique in ['summarize', 'keywords', 'bullet']:
            compressed = ExpandContext.compress_context(original_context, technique)
            compression_ratio = ExpandContext.count_tokens(compressed) / ExpandContext.count_tokens(original_context)
            print(f"\n{technique.upper()} COMPRESSION:")
            print("-" * 80)
            print(compressed)
            print("-" * 80)
            print(f"Compressed token count: {ExpandContext.count_tokens(compressed)}")
            print(f"Compression ratio: {compression_ratio:.2f} (lower is better)")

    @staticmethod
    def test_prune_context_layers(params:Dict = {}) -> None:
        pruning_test_prompt = "Write a tutorial on how to use pandas for data analysis."

        pruning_layers = {
            "role": "You are a data science instructor with 10+ years of experience teaching Python libraries.",
            "audience": "Your audience consists of beginner Python programmers who understand basic programming concepts but have no prior experience with data analysis.",
            "structure": "Structure the tutorial with these sections: Introduction, Installation, Loading Data, Basic Operations, Data Cleaning, Data Visualization, and a Practical Example.",
            "style": "Use a friendly, conversational tone. Include code snippets with comments explaining each line. Break down complex concepts into simple explanations.",
            "unnecessary": "Include details about the history of pandas and its development team. Mention that pandas was created by Wes McKinney in 2008 while he was at AQR Capital Management."
        }

        evaluation_criteria = [
            "Completeness - covers all essential concepts",
            "Clarity - explains concepts in an easy-to-understand way",
            "Code quality - provides useful, correct code examples",
            "Beginner-friendliness - assumes no prior knowledge of pandas",
            "Practicality - includes real-world applications"
        ]
        optimized_prompt, pruning_results = ExpandContext.prune_context_layers(pruning_test_prompt, pruning_layers, evaluation_criteria)
        print("\nOPTIMIZED PROMPT:")
        print("-" * 80)
        print(optimized_prompt)
        print("-" * 80)
        
        # Show quality scores for each configuration
        for config, data in pruning_results.items():
            print(f"{config}: Quality = {data['quality']:.2f}, Tokens = {data['tokens']}")

    ## 9. Context Expansion with Retrieval
    @staticmethod
    def test_context_expand_retrieval(params:Dict = {}) -> None:
        # Example knowledge base (in a real application, this would be much larger)
        sample_knowledge_base = [
            {
                "title": "Pandas Introduction",
                "content": "Pandas is a fast, powerful, flexible and easy to use open source data analysis and manipulation tool, built on top of the Python programming language. Key features include DataFrame objects, handling of missing data, and data alignment."
            },
            {
                "title": "Pandas Installation",
                "content": "To install pandas, run: pip install pandas. For Anaconda users, pandas comes pre-installed. You can import pandas with: import pandas as pd"
            },
            {
                "title": "Loading Data in Pandas",
                "content": "Pandas can read data from various sources including CSV, Excel, SQL databases, and JSON. Example: df = pd.read_csv('data.csv')"
            },
            {
                "title": "Data Cleaning with Pandas",
                "content": "Pandas provides functions for handling missing data, such as dropna() and fillna(). It also offers methods for removing duplicates and transforming data."
            },
            {
                "title": "Data Visualization with Pandas",
                "content": "Pandas integrates with matplotlib to provide plotting capabilities. Simple plots can be created with df.plot(). For more complex visualizations, use: import matplotlib.pyplot as plt"
            }
        ]
        # Test retrieval-augmented context expansion
        rag_test_prompt = "Write a brief tutorial on how to load data in pandas and handle missing values."
        rag_context = ExpandContext.create_rag_context(rag_test_prompt, "pandas loading data cleaning", sample_knowledge_base)
        print("RETRIEVAL-AUGMENTED CONTEXT:")
        print("-" * 80)
        print(rag_context)
        print("-" * 80)
        print(f"Token count: {ExpandContext.count_tokens(rag_context)}")
        # Generate response with RAG context
        rag_response, rag_latency = ExpandContext.generate_response(rag_context)
        print("\nRAG RESPONSE:")
        print("-" * 80)
        print(rag_response)
        print("-" * 80)